
// MemMapViewView.cpp : implementation of the CMemMapViewView class
//

#include "stdafx.h"
#include "MemMapView.h"

#include "MemMapViewDoc.h"
#include "MemMapViewView.h"
#include "MemoryMap.h"
#include "ProcessList.h"
#include "NtDll.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#endif

#define IDC_LIST 111

// CMemMapViewView

IMPLEMENT_DYNCREATE(CMemMapViewView, CView)

BEGIN_MESSAGE_MAP(CMemMapViewView, CView)
	ON_UPDATE_COMMAND_UI(ID_INDICATOR_SELECTION, OnMemorySelection)
	ON_WM_CREATE()
	ON_WM_ERASEBKGND()
	ON_WM_SIZE()
	ON_WM_SETFOCUS()
//	ON_NOTIFY(LVN_ITEMCHANGED, IDC_LIST, OnListItemChanged)
END_MESSAGE_MAP()

// CMemMapViewView construction/destruction

CMemMapViewView::CMemMapViewView() : m_List(this)
{
	m_MemSelected = 0;

}

CMemMapViewView::~CMemMapViewView()
{
}

BOOL CMemMapViewView::PreCreateWindow(CREATESTRUCT& cs)
{
	// TODO: Modify the Window class or styles here by modifying
	//  the CREATESTRUCT cs

	return CView::PreCreateWindow(cs);
}

// CMemMapViewView drawing

void CMemMapViewView::OnDraw(CDC* /*pDC*/)
{
	CMemMapViewDoc* pDoc = GetDocument();
	ASSERT_VALID(pDoc);
	if (!pDoc)
		return;

	// TODO: add draw code for native data here
}

void CMemMapViewView::OnRButtonUp(UINT nFlags, CPoint point)
{
	ClientToScreen(&point);
	OnContextMenu(this, point);
}

void CMemMapViewView::OnContextMenu(CWnd* pWnd, CPoint point)
{
	theApp.GetContextMenuManager()->ShowPopupMenu(IDR_POPUP_EDIT, point.x, point.y, this, TRUE);
}


// CMemMapViewView diagnostics

#ifdef _DEBUG
void CMemMapViewView::AssertValid() const
{
	CView::AssertValid();
}

void CMemMapViewView::Dump(CDumpContext& dc) const
{
	CView::Dump(dc);
}

CMemMapViewDoc* CMemMapViewView::GetDocument() const // non-debug version is inline
{
	ASSERT(m_pDocument->IsKindOf(RUNTIME_CLASS(CMemMapViewDoc)));
	return (CMemMapViewDoc*)m_pDocument;
}
#endif //_DEBUG


// CMemMapViewView message handlers

int CMemMapViewView::OnCreate(LPCREATESTRUCT lpCreateStruct)
{
	if (CView::OnCreate(lpCreateStruct) == -1)
		return -1;

	m_List.Create(WS_CHILD | WS_VISIBLE | LVS_REPORT | LVS_SHOWSELALWAYS, CRect(0, 0, 0, 0), this, IDC_LIST);
	m_List.SetExtendedStyle(LVS_EX_FULLROWSELECT);
	CFont font;
	font.CreatePointFont(110, _T("Courier New"));
	m_List.SetFont(&font);
	font.Detach();

	DWORD align = LVCFMT_CENTER;
	m_List.InsertColumn(0, _T(""), align, 0);
	m_List.InsertColumn(1, _T("Start Address"), align, 170);
	m_List.InsertColumn(2, _T("End Address"), align, 170);
	m_List.InsertColumn(3, _T("Size"), align, 170);
	m_List.InsertColumn(4, _T("State"), align, 110);
//	m_List.InsertColumn(5, _T("Allocation Protection"), align, 140);
	m_List.InsertColumn(5, _T("Protection"), align, 190);
	m_List.InsertColumn(6, _T("Type"), align, 120);
	m_List.InsertColumn(7, _T("Details"), LVCFMT_LEFT, 400);

	m_List.EnableMarkSortedColumn();

	return 0;
}

BOOL CMemMapViewView::OnEraseBkgnd(CDC* pDC)
{
	return TRUE;
}

void CMemMapViewView::OnSize(UINT nType, int cx, int cy)
{
	CView::OnSize(nType, cx, cy);

	if(m_List)
		m_List.MoveWindow(0, 0, cx, cy);
}

void CMemMapViewView::UpdateList() {
	m_List.DeleteAllItems();
	m_MemoryInfo.RemoveAll();
	m_List.LockWindowUpdate();

	CModuleList modules(GetDocument()->GetProcessID());
	CThreadList threads;

	HANDLE hTarget = ::OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, FALSE, GetDocument()->GetProcessID());

	CMemoryMap map;
	map.Init(GetDocument()->GetProcessID());
	MEMORY_BASIC_INFORMATION* info;
	CString addr;
	DWORD_PTR totalCommitted = 0, privateCommitted = 0;

	m_IsWow = true;
#ifdef _WIN64
	m_IsWow = CProcessList::IsWow64Process(GetDocument()->GetProcessID());
#endif
	while(info = map.GetNext()) {
		int n = m_List.InsertItem(m_List.GetItemCount(), _T(""));
		MemoryBuffer buffer;
		::ZeroMemory(&buffer, sizeof(buffer));
		*(MEMORY_BASIC_INFORMATION*)&buffer = *info;
		m_List.SetItemData(n, n);
		FormatAddress(addr, info->BaseAddress);
		m_List.SetItemText(n, 1, addr);
		FormatAddress(addr, (PBYTE)info->BaseAddress + info->RegionSize - 1);
		m_List.SetItemText(n, 2, addr);
		FormatAddress(addr, (PVOID)info->RegionSize);

		m_List.SetItemText(n, 3, addr);
		switch(info->State) {
			case MEM_FREE:
				addr = _T("Free");
				break;
			case MEM_COMMIT:
				addr = _T("Commited");
				break;
			case MEM_RESERVE:
				addr = _T("Reserved");
				break;
			default:
				addr = _T("N/A");
		}
		m_List.SetItemText(n, 4, addr);

		if(info->State != MEM_COMMIT) {
			addr = _T("N/A");
		}
		else {
			switch(info->Protect & 0xff) {
				case PAGE_EXECUTE: addr = _T("Execute"); break;
				case PAGE_EXECUTE_READ: addr = _T("Execute/Read"); break;
				case PAGE_EXECUTE_READWRITE: addr = _T("Execute/Read/Write"); break;
				case PAGE_NOACCESS: case 0: addr = _T("No Access"); break;
				case PAGE_READONLY: addr = _T("Read"); break;
				case PAGE_READWRITE: addr = _T("Read/Write"); break;
				case PAGE_WRITECOPY: addr = _T("Copy on Write"); break;
				case PAGE_EXECUTE_WRITECOPY: addr = _T("Execute/Copy on Write"); break;
			}

			switch(info->Protect & 0xff00) {
				case PAGE_GUARD: addr += _T(" + Guard"); break;
			}
		}

		m_List.SetItemText(n, 5, addr);

		if(info->State != MEM_FREE) {
			switch(info->Type) {
				case MEM_IMAGE: addr = _T("Image"); break;
				case MEM_MAPPED: addr = _T("Mapped"); break;
				case MEM_PRIVATE: addr = _T("Private"); break;
				default: addr = _T("N/A"); break;
			}

			if(info->Type == MEM_IMAGE) {
				modules.Reset();
				MODULEENTRY32* me;
				while(me = modules.GetNext()) {
					if((INT_PTR)info->BaseAddress >= (INT_PTR)me->modBaseAddr && (INT_PTR)info->BaseAddress < (INT_PTR)me->modBaseAddr + (INT_PTR)me->modBaseSize) {
						// found it!
						m_List.SetItemText(n, 7, me->szModule);
						::lstrcpy(buffer.szModule, me->szModule);
						break;
					}
				}
			}
			else if(info->Type == MEM_MAPPED && hTarget) {
				// get file name if possible...
				TCHAR path[MAX_PATH];
				if(::GetMappedFileName(hTarget, info->BaseAddress, path, MAX_PATH))
					m_List.SetItemText(n, 7, path);
			}
		}
		if(info->State == MEM_COMMIT) {
			totalCommitted += info->RegionSize;
			if(info->Type == MEM_PRIVATE)
				privateCommitted += info->RegionSize;
		}

		if(info->State != MEM_FREE && info->Type == MEM_PRIVATE) {
			if(hTarget) {
				THREADENTRY32* thinfo;
				CString buf;
				threads.Reset();
				while(thinfo = threads.GetNext()) {
					if(thinfo->th32OwnerProcessID != GetDocument()->GetProcessID()) continue;
					HANDLE hThread = ::OpenThread(THREAD_QUERY_INFORMATION, FALSE, thinfo->th32ThreadID);
					if(hThread) {
						THREAD_BASIC_INFORMATION tinfo;
						ULONG len;
						if(::NtQueryInformationThread(hThread, ThreadBasicInformation, &tinfo, sizeof(tinfo), &len) == 0) {
							NT_TIB tib;
							SIZE_T read;
							::ReadProcessMemory(hTarget, tinfo.TebBaseAddress, &tib, sizeof(tib), &read);
							if((INT_PTR)info->BaseAddress < (INT_PTR)tib.StackBase && (INT_PTR)info->BaseAddress >= (INT_PTR)tib.StackLimit) {
								buf.Format(_T("Thread %d (0x%X) Stack"), thinfo->th32ThreadID, thinfo->th32ThreadID);
								m_List.SetItemText(n, 7, buf);
							}
						}
						::CloseHandle(hThread);
					}
				}
			}
		}

		m_List.SetItemText(n, 6, addr);
		m_MemoryInfo.SetAt(n, buffer);
		::ZeroMemory(info, sizeof(*info));
	}
	if(hTarget)
		::CloseHandle(hTarget);

	GetDocument()->SetCommitSizes(privateCommitted, totalCommitted);

	m_List.UnlockWindowUpdate();
}

int CMemMapViewView::CMemMapList::OnCompareItems(LPARAM p1, LPARAM p2, int col) {
	MemoryBuffer& m1 = m_pView->m_MemoryInfo[p1], &m2 = m_pView->m_MemoryInfo[p2];
	switch(col) {
		case 1: case 2:
			return (INT_PTR)m1.BaseAddress - (INT_PTR)m2.BaseAddress > 0 ? 1 : -1;
		case 3:
			return (INT_PTR)m1.RegionSize - (INT_PTR)m2.RegionSize > 0 ? 1 : -1;
		case 4: 
			return (int)m1.State - (int)m2.State;
		case 5:
			return (INT)(m1.Protect & 0xff) - (INT)(m2.Protect & 0xff);
		case 6:
			return (int)m1.Type - (int)m2.Type;
		case 7:
			return ::lstrcmpi(m1.szModule, m2.szModule);

	}
	return 0;
}

void CMemMapViewView::OnUpdate(CView* /*pSender*/, LPARAM /*lHint*/, CObject* /*pHint*/)
{
	CWaitCursor wait;
	UpdateList();
}

void CMemMapViewView::OnSetFocus(CWnd* pOldWnd)
{
	CView::OnSetFocus(pOldWnd);

	m_List.SetFocus();
}

void CMemMapViewView::OnListItemChanged(NMHDR* pnmh, LRESULT* presult) {
}

void CMemMapViewView::OnMemorySelection(CCmdUI* pCmdUI) {
	DWORD_PTR selected = 0;
	for(POSITION pos = m_List.GetFirstSelectedItemPosition(); pos; ) {
		int n = (int)m_List.GetItemData(m_List.GetNextSelectedItem(pos));
		selected += m_MemoryInfo[n].RegionSize;
	}
	m_MemSelected = selected;
	CString buf;
	bool isMB = m_MemSelected > (DWORD_PTR)1 << 31;
	buf.Format(_T("Memory Selected: %d %cB"), isMB ? m_MemSelected >> 20 : m_MemSelected >> 10, isMB ? _T('M') : _T('K'));
	pCmdUI->Enable();
	pCmdUI->SetText(buf);
}

void CMemMapViewView::FormatAddress(CString& addr, PVOID p) {
	if(!m_IsWow)
		addr.Format(_T("%08X'%08X"), (DWORD)((DWORD_PTR)p >> 32), (DWORD)p);
	else
		addr.Format(_T("%08X"), p);
}
